package com.zjj.algorithm.learning.atguigu.tree.huffman;

import com.zjj.algorithm.learning.atguigu.tree.TreeNode;

import java.nio.charset.StandardCharsets;
import java.util.*;

/**
 * 使用 Huffman 编码压缩数据
 * @author zjj_admin
 * @date 2022/11/28 22:38
 */
public class HuffmanCode {

    public static void main(String[] args) {
        String str = "I like like java and like shu ju jie gou";
//        String str = "i like like like java do you like a java";
        byte[] strBytes = str.getBytes(StandardCharsets.UTF_8);
        int length = strBytes.length;
        System.out.println("length = " + length);
        List<CodeNode> nodes = getNode(strBytes);
        System.out.println("nodes = " + nodes);

        CodeNode codeNode = createHuffmanTree(nodes);
        System.out.println("codeNode = " + codeNode);

        getCode(codeNode,"",new StringBuilder());
        System.out.println("huffmanCode = " + huffmanCode);

        StringBuilder res = new StringBuilder();
        for (int i = 0; i < strBytes.length; i++) {
            String code = huffmanCode.get(strBytes[i]);
            res.append(code);
        }
        System.out.println("res = " + res.toString());
        System.out.println("res 的长度 = " + res.toString().length());

        StringBuilder res2 = new StringBuilder();
        //将数据还原
        int i = 0;
        int j = 1;
        while (true){
            String sub = res.substring(i, j + i);
            if(huffmanCode.containsValue(sub)){
                for (Map.Entry<Byte, String> entry : huffmanCode.entrySet()) {
                    if(entry.getValue().equals(sub)){
                        String data = new String(new byte[]{entry.getKey()}, StandardCharsets.UTF_8);
                        res2.append(data);
                        i = i + j;
                        j = 1;

                    }
                }
            }else {
                j ++;
            }
            if(i == res.length()){
                break;
            }
        }
        System.out.println("res2 = " + res2);



    }



    private static Map<Byte,String> huffmanCode = new HashMap<>();

    /**
     * 对每一个叶子节点生成 Huffman 编码
     * @param node
     * @param code
     * @param stringBuilder
     */
    private static void getCode(CodeNode node,String code,StringBuilder stringBuilder){
        StringBuilder stringBuilder2 = new StringBuilder(stringBuilder);
        stringBuilder2.append(code);
        if (node != null) {
            if (node.getData() == null) {
                getCode(node.getLeft(), "0", stringBuilder2);
                getCode(node.getRight(), "1", stringBuilder2);
            }else {
                huffmanCode.put(node.getData(), stringBuilder2.toString());
            }
        }
    }


    /**
     * 根据字节数组返回 List<CodeNode>
     * @param bytes
     * @return
     */
    private static List<CodeNode> getNode(byte[] bytes){
        List<CodeNode> nodes = new ArrayList<>();
        Map<Byte,Integer> counts = new HashMap<>();
        for (byte b : bytes) {
            Integer count = counts.get(b);
            if(null == count){
                counts.put(b,1);
            }else {
                counts.put(b,count + 1);
            }
        }
        //遍历每一个键值，将其存储到 nodes 中
        for (Map.Entry<Byte, Integer> entry : counts.entrySet()) {
            nodes.add(new CodeNode(entry.getKey(), entry.getValue()));
        }
        return nodes;
    }

    /**
     * 构建哈夫曼树
     * @param nodes
     * @return
     */
    private static CodeNode createHuffmanTree(List<CodeNode> nodes){
        while (nodes.size() > 1){
            //对节点进行排序
            Collections.sort(nodes);
            //获取最小的两个节点
            CodeNode leftNode = nodes.get(0);
            CodeNode rightNode = nodes.get(1);
            //构建新的节点
            CodeNode parentNode = new CodeNode(null, leftNode.getWeight() + rightNode.getWeight());
            parentNode.setLeft(leftNode);
            parentNode.setRight(rightNode);

            nodes.remove(leftNode);
            nodes.remove(rightNode);

            nodes.add(parentNode);
        }
        return nodes.get(0);
    }




}



